use criterion::{black_box, criterion_group, criterion_main, Criterion};
use std::sync::Arc;
use std::time::Duration;
use tokio::runtime::Runtime;
use anyhow::Result;

use protoactor_rs::actor::{Actor, ActorContext, ActorError, PID, Props};
use protoactor_rs::system::ActorSystem;
use protoactor_rs::remote::RemoteActorSystem;

// PingActor - 用于测试本地和远程消息传递性能
struct PingActor {
    count: usize,
    pongs: usize,
    batch_size: usize,
}

impl PingActor {
    fn new(count: usize, batch_size: usize) -> Self {
        Self {
            count,
            pongs: 0,
            batch_size,
        }
    }
}

#[async_trait::async_trait]
impl Actor for PingActor {
    type Context = protoactor_rs::actor::DefaultActorContext;
    
    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn std::any::Any + Send>) -> Result<(), ActorError> {
        if let Some(_) = msg.downcast_ref::<Pong>() {
            let mut this = self.clone();
            this.pongs += 1;
            
            if this.pongs >= this.count {
                ctx.stop();
            } else if this.pongs % this.batch_size == 0 {
                // 发送下一批ping消息
                for _ in 0..this.batch_size {
                    ctx.send(&ctx.sender().unwrap(), Ping {}).await?;
                }
            }
        }
        Ok(())
    }
}

impl Clone for PingActor {
    fn clone(&self) -> Self {
        Self {
            count: self.count,
            pongs: self.pongs,
            batch_size: self.batch_size,
        }
    }
}

// PongActor - 简单地回复Ping消息
struct PongActor;

#[async_trait::async_trait]
impl Actor for PongActor {
    type Context = protoactor_rs::actor::DefaultActorContext;
    
    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn std::any::Any + Send>) -> Result<(), ActorError> {
        if let Some(_) = msg.downcast_ref::<Ping>() {
            ctx.respond(Pong {});
        }
        Ok(())
    }
}

#[derive(Debug)]
struct Ping;

#[derive(Debug)]
struct Pong;

// 本地PingPong基准测试
fn bench_local_ping_pong(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    
    c.bench_function("local_ping_pong", |b| {
        b.iter(|| {
            rt.block_on(async {
                let system = Arc::new(ActorSystem::new().unwrap());
                
                // 创建PongActor
                let pong_props = Props::new(Box::new(|| Box::new(PongActor)));
                let pong_pid = system.spawn(pong_props).await.unwrap();
                
                // 创建PingActor
                let ping_props = Props::new(Box::new(|| {
                    Box::new(PingActor::new(1000, 100))
                }));
                let ping_pid = system.spawn(ping_props).await.unwrap();
                
                // 发送初始消息
                for _ in 0..100 {
                    system.send(&pong_pid, Ping {}).await.unwrap();
                }
                
                // 等待完成
                tokio::time::sleep(Duration::from_millis(100)).await;
                
                // 停止actors
                system.stop(&ping_pid).await.unwrap();
                system.stop(&pong_pid).await.unwrap();
            });
        })
    });
}

// 远程PingPong基准测试
fn bench_remote_ping_pong(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    
    c.bench_function("remote_ping_pong", |b| {
        b.iter(|| {
            rt.block_on(async {
                // 创建两个远程Actor系统
                let system1 = Arc::new(RemoteActorSystem::new("127.0.0.1:8081").unwrap());
                let system2 = Arc::new(RemoteActorSystem::new("127.0.0.1:8082").unwrap());
                
                // 启动系统
                system1.start().await.unwrap();
                system2.start().await.unwrap();
                
                // 在system1上创建PongActor
                let pong_props = Props::new(Box::new(|| Box::new(PongActor)));
                let pong_pid = system1.spawn_remote("127.0.0.1:8081", "pong", "", Duration::from_secs(5)).await.unwrap();
                
                // 在system2上创建PingActor
                let ping_props = Props::new(Box::new(|| {
                    Box::new(PingActor::new(1000, 100))
                }));
                let ping_pid = system2.spawn_remote("127.0.0.1:8082", "ping", "", Duration::from_secs(5)).await.unwrap();
                
                // 发送初始消息
                for _ in 0..100 {
                    system2.send(&pong_pid, Ping {}).await.unwrap();
                }
                
                // 等待完成
                tokio::time::sleep(Duration::from_secs(1)).await;
                
                // 停止actors
                system1.stop().await.unwrap();
                system2.stop().await.unwrap();
                
                // 停止系统
                system1.stop().await.unwrap();
                system2.stop().await.unwrap();
            });
        })
    });
}

// SkyNet基准测试 - 创建大量Actor并测试消息传递
fn bench_skynet(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    
    c.bench_function("skynet", |b| {
        b.iter(|| {
            rt.block_on(async {
                let system = Arc::new(ActorSystem::new().unwrap());
                
                // 创建Actor树
                let root_props = Props::new(Box::new(|| {
                    Box::new(SkyNetActor::new(1000, 10))
                }));
                let root_pid = system.spawn(root_props).await.unwrap();
                
                // 发送开始消息
                system.send(&root_pid, Start {}).await.unwrap();
                
                // 等待完成
                tokio::time::sleep(Duration::from_secs(1)).await;
                
                // 停止root actor
                system.stop(&root_pid).await.unwrap();
            });
        })
    });
}

// SkyNet Actor
struct SkyNetActor {
    div: usize,
    size: usize,
    count: usize,
    children: Vec<PID>,
}

impl SkyNetActor {
    fn new(size: usize, div: usize) -> Self {
        Self {
            div,
            size,
            count: 0,
            children: Vec::new(),
        }
    }
}

#[derive(Debug)]
struct Start;

#[derive(Debug)]
struct Message(usize);

#[async_trait::async_trait]
impl Actor for SkyNetActor {
    type Context = protoactor_rs::actor::DefaultActorContext;
    
    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn std::any::Any + Send>) -> Result<(), ActorError> {
        if let Some(_) = msg.downcast_ref::<Start>() {
            let mut this = self.clone();
            if this.size <= this.div {
                return Ok(());
            }
            
            let new_size = this.size / this.div;
            for _ in 0..this.div {
                let child_props = Props::new(Box::new(move || {
                    Box::new(SkyNetActor::new(new_size, this.div))
                }));
                let child_pid = ctx.spawn(child_props).await?;
                this.children.push(child_pid);
                
                // 向子Actor发送开始消息
                ctx.send(&child_pid, Start {}).await?;
            }
        } else if let Some(Message(num)) = msg.downcast_ref::<Message>() {
            let mut this = self.clone();
            this.count += num;
            
            if this.children.is_empty() {
                ctx.respond(Message(1));
            } else {
                for child in &this.children {
                    ctx.send(child, Message(*num)).await?;
                }
            }
        }
        Ok(())
    }
}

impl Clone for SkyNetActor {
    fn clone(&self) -> Self {
        Self {
            div: self.div,
            size: self.size,
            count: self.count,
            children: self.children.clone(),
        }
    }
}

criterion_group!(benches, bench_local_ping_pong, bench_remote_ping_pong, bench_skynet);
criterion_main!(benches); 